from torch_geometric.loader import LinkNeighborLoader, ZipLoader
import torch
import copy
import torch_geometric.transforms as T
from data_package_003 import my_dataset
from torch_geometric.loader import DataLoader
from torch_geometric.data import Data
from torch_geometric.data import HeteroData


def descrition_data(_data_dict, _data):
    _descrition_data = {}
    all_list = []
    vdata = _data.edge_label_index.t().tolist()
    # print(vdata)
    for i2j in vdata:
        all_list.append((i2j, _data_dict.all_dict[i2j[0]], _data_dict.all_dict[i2j[1]]))
    return all_list


# <class 'torch_geometric.data.hetero_data.HeteroData'>
# <class 'torch_geometric.data.hetero_data.HeteroData'>
train_data_list = []

#
# print(next(iter(my_dataset)))
datas = []

aoh_list = []
atb_list = []
host_list = []
host_node_id_list = []
alarm_list = []
alarm_node_id_list = []
btree_list = []
btree_node_id_list = []
"""
Data(edge_index=[2, 1136], node_id=[389], x=[389, 768], node_type=[389], edge_type=[1136])
Data(edge_index=[2, 948], node_id=[389], x=[389, 768], node_type=[389], edge_type=[948])
Data(edge_index=[2, 1124], node_id=[389], x=[389, 768], node_type=[389], edge_type=[1124])
Data(edge_index=[2, 1170], node_id=[367], x=[367, 768], node_type=[367], edge_type=[1170])
"""

first = True
for batch in my_dataset:
    if first == True:
        first = False
        host_list.append(batch["host"].x)
        btree_list.append(batch["btree"].x)
        host_node_id_list.append(batch["host"].node_id)
        btree_node_id_list.append(batch["btree"].node_id)

    alarm_list.append(batch["alarm"].x)
    alarm_node_id_list.append(batch["alarm"].node_id)

    aoh_list.append(batch[("alarm", "on", "host")])
    atb_list.append(batch[("alarm", "to", "btree")])

data = HeteroData()
data["btree"].node_id = torch.IntTensor(torch.cat(btree_node_id_list))
data['btree'].x = torch.cat(btree_list)
# data["host"].node_id = torch.IntTensor(torch.cat(host_node_id_list))
# data['host'].x = torch.cat(host_list)
data["alarm"].node_id = torch.IntTensor(torch.cat(alarm_node_id_list))
data['alarm'].x = torch.cat(alarm_list, dim=0)
# data[('alarm', 'on', 'host')].edge_index = torch.cat([i["edge_index"] for i in aoh_list],dim=1)
data[('alarm', 'to', 'btree')].edge_index = torch.cat([i["edge_index"] for i in atb_list], dim=1)

# print("data1",data)
# del data['alarm'].num_nodes
# print("data2",data)

data = T.ToUndirected()(data)
# print("1",data)
# del data['btree', 'rev_to', 'alarm'].edge_label
print("2", data)
print(data.metadata()[1])
from torch_geometric.nn import HeteroConv, SAGEConv
import torch.nn.functional as F

metadata = data.metadata()
hidden_channels = 512
conv = HeteroConv({
    edge_type: SAGEConv((-1, -1), hidden_channels)
    for edge_type in metadata[1]
})
print(conv.convs,"$$$$$$$$$$")
"""
ModuleDict(
  (<alarm___to___btree>): SAGEConv((-1, -1), 512, aggr=mean)
  (<btree___rev_to___alarm>): SAGEConv((-1, -1), 512, aggr=mean)
) 
"""
#####################
from fasttext import FTModel
import torch
from create_alarms_data_002 import ORIGIN_DATA


def get_sentence_vectors(texts):
    word_embeddings_list = []
    for text in texts:
        encoded_input = FTModel.embedding_text_to_tensor(text)
        word_embeddings_list.append(encoded_input.unsqueeze(0))
    return torch.cat(word_embeddings_list, dim=0)


device = "cuda:0" if torch.cuda.is_available() else "cpu"


def translate_batch(batch):
    _btree_list = [ORIGIN_DATA.btrees_origin_infos[id] for id in batch["btree"].node_id.tolist()]
    _alarm_list = [ORIGIN_DATA.alarms_origin_infos[id] for id in batch["alarm"].node_id.tolist()]
    batch["btree"].x = get_sentence_vectors(_btree_list)
    # batch["host"].x = get_sentence_vectors(_host_list)
    batch["alarm"].x = get_sentence_vectors(_alarm_list)
    _batch = batch.to(device)
    return _batch


data = translate_batch(data)
print(data)
print("#################")
# 实际上已经开始进行分层计算
for conv in [conv]:
    x_dict = conv(data.x_dict, data.edge_index_dict)
    x_dict = {key: F.leaky_relu(x) for key, x in x_dict.items()}
print(x_dict["alarm"].shape)
print(x_dict["btree"].shape)
print("#################")

# 定义元路径
from torch_geometric.transforms import AddMetaPaths

metapaths = [[('alarm', 'btree'), ("btree", "alarm")]]
# 创建一个转换对象，用于添加元路径
# drop_orig_edge_types 去掉之前的连接边信息，drop_unconnected_node_types 去掉不参与连接的node信息
transform = AddMetaPaths(metapaths=metapaths, drop_unconnected_node_types=True)
transform = AddMetaPaths(metapaths=metapaths, drop_orig_edge_types=True)
transform = AddMetaPaths(metapaths=metapaths, drop_orig_edge_types=True, drop_unconnected_node_types=True)
data = transform(data)
print("3", data)
"""
2 HeteroData(
  btree={
    node_id=[11],
    x=[11],
  },
  alarm={
    node_id=[746],
    x=[746],
  },
  (alarm, to, btree)={ edge_index=[2, 1146] },
  (btree, rev_to, alarm)={ edge_index=[2, 1146] }
)
3 HeteroData(
  metapath_dict={ (alarm, metapath_0, alarm)=[2] },
  alarm={
    node_id=[746],
    x=[746],
  },
  (alarm, metapath_0, alarm)={ edge_index=[2, 111074] }
)
"""
